package com.day;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.google.gson.Gson;
import okhttp3.*;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.io.IOException;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.text.SimpleDateFormat;
import java.util.*;

/**
 * 大模型WebSocket通信类，用于与大模型服务进行对话。
 */
public class BigModelNew extends WebSocketListener {
    // 大模型服务的基础URL和鉴权信息
    public static final String hostUrl = "https://spark-api.xf-yun.com/v3.5/chat";
    public static final String appid = "fc36b5c2";
    public static final String apiSecret = "ZjRlNWM0NjE3MWVhMjA3MmE2ZGZiYTk5";
    public static final String apiKey = "052929db22adf23deb5a56ce1e0a7bac";

    // 用于存储对话历史
    public static List<RoleContent> historyList = new ArrayList<>();
    // 存储大模型给出的总答案
    public static String totalAnswer = "";
    // 当前待回答的问题
    public static String NewQuestion = "";

    // Gson实例用于JSON解析
    public static final Gson gson = new Gson();

    // 用户ID和WebSocket关闭标志
    private String userId;
    private Boolean wsCloseFlag;

    // 控制是否提示用户输入的标志
    private static Boolean totalFlag = true;

    /**
     * 构造函数
     *
     * @param userId 用户ID
     * @param wsCloseFlag WebSocket关闭标志
     */
    public BigModelNew(String userId, Boolean wsCloseFlag) {
        this.userId = userId;
        this.wsCloseFlag = wsCloseFlag;
    }

    /**
     * 主函数，用于启动程序并发起WebSocket连接。
     */
    public static void main(String[] args) throws Exception {
        while (true) {
            if (totalFlag) {
                // 读取用户输入
                Scanner scanner = new Scanner(System.in);
                System.out.print("我：");
                totalFlag = false;
                NewQuestion = scanner.nextLine();

                // 构建鉴权URL并发起WebSocket连接
                String authUrl = getAuthUrl(hostUrl, apiKey, apiSecret);
                OkHttpClient client = new OkHttpClient.Builder().build();
                String url = authUrl.toString().replace("http://", "ws://").replace("https://", "wss://");
                Request request = new Request.Builder().url(url).build();
                for (int i = 0; i < 1; i++) {
                    totalAnswer = "";
                    WebSocket webSocket = client.newWebSocket(request, new BigModelNew(i + "", false));
                }
            } else {
                // 等待一段时间后再次提示用户输入
                Thread.sleep(200);
            }
        }
    }

    /**
     * 判断是否可以添加历史记录，避免历史记录过多。
     *
     * @return 返回true表示可以添加历史记录，否则返回false。
     */
    public static boolean canAddHistory() {
        int history_length = 0;
        for (RoleContent temp : historyList) {
            history_length += temp.content.length();
        }
        // 如果历史记录总长度超过1.2万字符，移除部分最早的历史记录并返回false，否则返回true
        if (history_length > 12000) {
            historyList.subList(0, 5).clear(); // 移除前5条记录
            return false;
        } else {
            return true;
        }
    }

    /**
     * 用于发送音频与参数的线程类。
     */
    class MyThread extends Thread {
        private WebSocket webSocket;

        public MyThread(WebSocket webSocket) {
            this.webSocket = webSocket;
        }

        public void run() {
            try {
                // 构建发送给大模型的请求JSON
                JSONObject requestJson = new JSONObject();
                JSONObject header = new JSONObject();
                header.put("app_id", appid);
                header.put("uid", UUID.randomUUID().toString().substring(0, 10));

                JSONObject parameter = new JSONObject();
                JSONObject chat = new JSONObject();
                chat.put("domain", "generalv2");
                chat.put("temperature", 0.5);
                chat.put("max_tokens", 4096);
                parameter.put("chat", chat);

                JSONObject payload = new JSONObject();
                JSONObject message = new JSONObject();
                JSONArray text = new JSONArray();

                // 添加历史问题到请求
                if (!historyList.isEmpty()) {
                    for (RoleContent tempRoleContent : historyList) {
                        text.add(JSON.toJSON(tempRoleContent));
                    }
                }

                // 添加当前问题到请求
                RoleContent roleContent = new RoleContent();
                roleContent.role = "user";
                roleContent.content = NewQuestion;
                text.add(JSON.toJSON(roleContent));
                historyList.add(roleContent); // 添加到历史记录

                message.put("text", text);
                payload.put("message", message);
                requestJson.put("header", header);
                requestJson.put("parameter", parameter);
                requestJson.put("payload", payload);

                // 发送请求
                webSocket.send(requestJson.toString());
                // 等待服务端返回
                while (true) {
                    Thread.sleep(200);
                    if (wsCloseFlag) {
                        break;
                    }
                }
                webSocket.close(1000, "");
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * WebSocket打开时的回调方法。
     */
    @Override
    public void onOpen(WebSocket webSocket, Response response) {
        super.onOpen(webSocket, response);
        System.out.print("大模型：");
        MyThread myThread = new MyThread(webSocket);
        myThread.start();
    }

    /**
     * 接收到WebSocket消息时的回调方法。
     */
    @Override
    public void onMessage(WebSocket webSocket, String text) {
        JsonParse myJsonParse = gson.fromJson(text, JsonParse.class);
        if (myJsonParse.header.code != 0) {
            System.out.println("发生错误，错误码为：" + myJsonParse.header.code);
            System.out.println("本次请求的sid为：" + myJsonParse.header.sid);
            webSocket.close(1000, "");
        }
        List<Text> textList = myJsonParse.payload.choices.text;
        for (Text temp : textList) {
            System.out.print(temp.content);
            totalAnswer += temp.content;
        }
        if (myJsonParse.header.status == 2) {
            // 更新历史记录并关闭WebSocket连接
            if (canAddHistory()) {
                RoleContent roleContent = new RoleContent();
                roleContent.setRole("assistant");
                roleContent.setContent(totalAnswer);
                historyList.add(roleContent);
            } else {
                historyList.remove(0);
                RoleContent roleContent = new RoleContent();
                roleContent.setRole("assistant");
                roleContent.setContent(totalAnswer);
                historyList.add(roleContent);
            }
            wsCloseFlag = true;
            totalFlag = true;
        }
    }

    /**
     * WebSocket失败时的回调方法。
     */
    @Override
    public void onFailure(WebSocket webSocket, Throwable t, Response response) {
        super.onFailure(webSocket, t, response);
        try {
            if (null != response) {
                int code = response.code();
                System.out.println("onFailure code:" + code);
                System.out.println("onFailure body:" + response.body().string());
                if (101 != code) {
                    System.out.println("connection failed");
                    System.exit(0);
                }
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 鉴权方法，用于生成访问大模型服务所需的授权信息。
     *
     * @param hostUrl    大模型服务的URL
     * @param apiKey     API密钥
     * @param apiSecret  API密钥 secret
     * @return 鉴权URL
     * @throws Exception 加密或构建URL过程中可能出现的异常
     */
    public static String getAuthUrl(String hostUrl, String apiKey, String apiSecret) throws Exception {
        URL url = new URL(hostUrl);
        SimpleDateFormat format = new SimpleDateFormat("EEE, dd MMM yyyy HH:mm:ss z", Locale.US);
        format.setTimeZone(TimeZone.getTimeZone("GMT"));
        String date = format.format(new Date());
        String preStr = "host: " + url.getHost() + "\n" +
                "date: " + date + "\n" +
                "GET " + url.getPath() + " HTTP/1.1";
        // 使用SHA256和HMAC算法对预字符串进行加密
        Mac mac = Mac.getInstance("hmacsha256");
        SecretKeySpec spec = new SecretKeySpec(apiSecret.getBytes(StandardCharsets.UTF_8), "hmacsha256");
        mac.init(spec);
        byte[] hexDigits = mac.doFinal(preStr.getBytes(StandardCharsets.UTF_8));
        // 使用Base64对加密结果进行编码
        String sha = Base64.getEncoder().encodeToString(hexDigits);
        String authorization = String.format("api_key=\"%s\", algorithm=\"%s\", headers=\"%s\", signature=\"%s\"", apiKey, "hmac-sha256", "host date request-line", sha);
        HttpUrl httpUrl = Objects.requireNonNull(HttpUrl.parse("https://" + url.getHost() + url.getPath())).newBuilder().//
                addQueryParameter("authorization", Base64.getEncoder().encodeToString(authorization.getBytes(StandardCharsets.UTF_8))).//
                addQueryParameter("date", date).//
                addQueryParameter("host", url.getHost()).//
                build();
        return httpUrl.toString();
    }

    // JSON解析相关的内部类
    class JsonParse {
        Header header;
        Payload payload;
    }

    class Header {
        int code;
        int status;
        String sid;
    }

    class Payload {
        Choices choices;
    }

    class Choices {
        List<Text> text;
    }

    class Text {
        String role;
        String content;
    }

    class RoleContent {
        String role;
        String content;

        public String getRole() {
            return role;
        }

        public void setRole(String role) {
            this.role = role;
        }

        public String getContent() {
            return content;
        }

        public void setContent(String content) {
            this.content = content;
        }
    }
}
